import torch
import torchvision 
from tqdm import tqdm
#直接引入matplotlib包是不能引入pyplot模块的
#import matplotlib 会出错
import matplotlib.pyplot as plt
import matplotlib as mpl


#By: Elwin https://editor.csdn.net/md?not_checkout=1&articleId=112980305

class Net(torch.nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.model = torch.nn.Sequential(
            #The size of the picture is 28x28
            torch.nn.Conv2d(in_channels = 1,out_channels = 16,kernel_size = 3,stride = 1,padding = 1),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size = 2,stride = 2),
            
            #The size of the picture is 14x14
            torch.nn.Conv2d(in_channels = 16,out_channels = 32,kernel_size = 3,stride = 1,padding = 1),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size = 2,stride = 2),
            
            #The size of the picture is 7x7
            torch.nn.Conv2d(in_channels = 32,out_channels = 64,kernel_size = 3,stride = 1,padding = 1),
            torch.nn.ReLU(),
            
            torch.nn.Flatten(),
            torch.nn.Linear(in_features = 7 * 7 * 64,out_features = 128),
            torch.nn.ReLU(),
            torch.nn.Linear(in_features = 128,out_features = 10),
            torch.nn.Softmax(dim=1)
        )
        
    def forward(self,input):
        output = self.model(input)
        return output

def train():
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor(), torchvision.transforms.Normalize(mean = [0.5],std = [0.5])])
        
    BATCH_SIZE = 256
    EPOCHS = 10
    trainData = torchvision.datasets.MNIST('./data/', train = True, transform = transform,download = False)
    testData = torchvision.datasets.MNIST('./data/', train = False, transform = transform)


    trainDataLoader = torch.utils.data.DataLoader(dataset = trainData,batch_size = BATCH_SIZE,shuffle = True)
    testDataLoader = torch.utils.data.DataLoader(dataset = testData,batch_size = BATCH_SIZE)
    net = Net()
    print(net.to(device))

    lossF = torch.nn.CrossEntropyLoss() 
    optimizer = torch.optim.Adam(net.parameters())

    history = {'Test Loss':[],'Test Accuracy':[]}
    for epoch in range(1,EPOCHS + 1):
        processBar = tqdm(trainDataLoader,unit = 'step')
        net.train(True)
        for step,(trainImgs,labels) in enumerate(processBar):
            trainImgs = trainImgs.to(device)
            labels = labels.to(device)

            net.zero_grad()
            outputs = net(trainImgs)
            loss = lossF(outputs,labels)
            predictions = torch.argmax(outputs, dim = 1)
            accuracy = torch.sum(predictions == labels)/labels.shape[0]
            loss.backward()

            optimizer.step()
            processBar.set_description("[%d/%d] Loss: %.4f, Acc: %.4f" % 
                                    (epoch,EPOCHS,loss.item(),accuracy.item()))
            
            if step == len(processBar)-1:
                correct,totalLoss = 0,0
                net.train(False)
                with torch.no_grad():
                    for testImgs,labels in testDataLoader:
                        testImgs = testImgs.to(device)
                        labels = labels.to(device)
                        outputs = net(testImgs)
                        loss = lossF(outputs,labels)
                        predictions = torch.argmax(outputs,dim = 1)

                        totalLoss += loss
                        correct += torch.sum(predictions == labels)

                        testAccuracy = correct/(BATCH_SIZE * len(testDataLoader))
                        testLoss = totalLoss/len(testDataLoader)
                        history['Test Loss'].append(testLoss.item())
                        history['Test Accuracy'].append(testAccuracy.item())
                
                processBar.set_description("[%d/%d] Loss: %.4f, Acc: %.4f, Test Loss: %.4f, Test Acc: %.4f" % 
                                    (epoch,EPOCHS,loss.item(),accuracy.item(),testLoss.item(),testAccuracy.item()))
        processBar.close()


    plt.plot(history['Test Loss'],label = 'Test Loss')
    plt.legend(loc='best')
    plt.grid(True)
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.show()

    plt.plot(history['Test Accuracy'],color = 'red',label = 'Test Accuracy')
    plt.legend(loc='best')
    plt.grid(True)
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.show()

    # 保存整个模型
    torch.save(net,'./model.pth')
    # 只保存权重
    # torch.save(model.state_dict(), 'save.pt')


def predict():
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    model = torch.load('./model.pth', map_location=torch.device('cpu'))
    torch.no_grad()

    transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor(),
                                    torchvision.transforms.Normalize(mean = [0.5],std = [0.5])])
    BATCH_SIZE = 256
    EPOCHS = 10
    testData = torchvision.datasets.MNIST('./data/', train = False, transform = transform)
    testDataLoader = torch.utils.data.DataLoader(dataset = testData,batch_size = BATCH_SIZE)
    for testImgs,labels in testDataLoader:
        #testImgs = testImgs.to(device)
        #labels = labels.to(device)
        outputs = model(testImgs)
        predictions = torch.argmax(outputs,dim = 1)
        print(predictions)
        print(labels)
        correct = torch.sum(predictions == labels)
        print(correct)
        return


if __name__ == '__main__':
    predict()
    #train()

